import datetime
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import pandas as pd
import basic_func.general_func as bf
  
def Gantt_plot():
  colors = bf.color_selector(df,params)
  fig,ax = plt.subplots(figsize = (10,6),layout = 'tight')
  pos_start = 6
  step = 0.5
  count = len(df)
  for i in range(count):
    start = mdates.date2num(df['start'].iloc[i])
    end = mdates.date2num(df['end'].iloc[i])
    bottom = 0.5*(count-i-1) + pos_start
    width = end-start
    plt.barh(bottom,width,left = start,
             height = 0.3,
             label = df.index[i],
             color = colors[i])
  rule = mdates.rrulewrapper(mdates.DAILY,interval=20)
  loc = mdates.RRuleLocator(rule)
  formatter = mdates.DateFormatter("%m月%d日")

  ax = plt.gca() # 获取当前的坐标轴
  ax.xaxis.set_major_locator(loc)
  ax.xaxis.set_major_formatter(formatter)
  ax.set_yticks(np.arange(pos_start,
                          pos_start+count*step,
                          step),
                labels = df.index,
                fontsize = 15)

  plt.xticks(rotation=params['xlabel_rotation'],ha = 'right',fontsize = 13)
  plt.title(params['title'],
            fontsize= params['title_fontsize'],
            fontweight = 'bold')
  plt.grid(params['grid_show'],
           axis = params['grid_axis'],
           color = params['grid_color'],
           ls = params['grid_ls'],
           alpha = params['grid_alpha'])
  plt.show()


if __name__=='__main__':
  df,params = bf.read_data('test_data/项目进度甘特图.xlsx')
  df = df[::-1]
  df['start'] = pd.to_datetime(df.iloc[:,0])
  df['end'] = pd.to_datetime(df.iloc[:,1])
  bf.zh_font(params)
  Gantt_plot()
